import os
import json
import sys
import re
from datetime import datetime

# 使用input方法获取模型目录名
def get_model_dir_from_input():
    print("请输入模型目录名（位于runs/detect/下的子目录）")
    print("例如: 11s_lowlr_bigend_20251123_125057")
    model_dir = input("模型目录名: ").strip()
    
    if not model_dir:
        print("错误: 模型目录名不能为空")
        # 递归调用直到用户输入有效内容
        return get_model_dir_from_input()
    
    return model_dir

# 构建training_configs.json的完整路径
def build_config_path(model_dir):
    # 假设模型目录位于runs/detect/下
    base_dir = os.path.join("runs", "detect", model_dir, "weights")
    config_path = os.path.join(base_dir, "training_configs.json")
    return config_path

# 读取training_configs.json文件
def read_training_configs(config_path):
    if not os.path.exists(config_path):
        print(f"警告: 找不到配置文件 {config_path}")
        # 返回空字典而不是退出程序，以便后续错误处理
        return {}
    
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            configs = json.load(f)
        return configs
    except Exception as e:
        print(f"警告: 读取配置文件时出错: {str(e)}")
        # 返回空字典而不是退出程序
        return {}

# 解析training_configs.json，提取MS_CSV, MS_NML, MS_AVG, MS_END参数
def extract_ms_configs(configs):
    # 提取各个阶段的配置
    ms_configs = {}
    
    # 查找包含MS_CSV, MS_NML, MS_AVG, MS_END的配置
    for key, value in configs.items():
        if isinstance(value, dict):
            if 'hsv_h' in value and 'hsv_s' in value and 'hsv_v' in value:
                # 这可能是一个MS配置
                if 'CSV' in key.upper():
                    ms_configs['MS_CSV'] = value
                elif 'NML' in key.upper():
                    ms_configs['MS_NML'] = value
                elif 'AVG' in key.upper():
                    ms_configs['MS_AVG'] = value
                elif 'END' in key.upper():
                    ms_configs['MS_END'] = value
    
    # 如果没有找到，尝试直接从顶层提取
    if not ms_configs:
        # 检查是否有各个阶段的配置
        phases = ['MS_CSV', 'MS_NML', 'MS_AVG', 'MS_END']
        for phase in phases:
            if phase.lower() in configs:
                ms_configs[phase] = configs[phase.lower()]
            elif phase in configs:
                ms_configs[phase] = configs[phase]
    
    return ms_configs

# 将Python字典转换为格式化的字符串
def dict_to_formatted_string(config_dict, indent=8):
    lines = []
    lines.append('{')
    for key, value in config_dict.items():
        # 添加注释
        comment = ''
        # if key == 'patience':
        #     comment = '     # 早停耐心值'
        if key in ['hsv_h', 'hsv_s', 'hsv_v']:
            comment = '     # 色彩增强'
        elif key in ['degrees', 'translate', 'scale', 'shear', 'perspective']:
            comment = '     # 几何变换'
        elif key in ['flipud', 'fliplr']:
            comment = '     # 翻转增强'
        elif key in ['mosaic', 'mixup', 'cutmix', 'copy_paste', 'erasing']:
            comment = '     # 图像裁切'
        elif key == 'lr0':
            comment = '     # 初始学习率'
        elif key == 'cos_lr':
            comment = '     # 余弦学习率调度'
        # elif key == 'warmup_epochs':
        #     comment = '     # 预热轮次'
        elif key == 'dropout':
            comment = '     #  dropout比例'
        elif key == 'weight_decay':
            comment = '     # 权重衰减'
        elif key == 'label_smoothing':
            comment = '     # 标签平滑'
        elif key == 'batch':
            comment = '     # 批次大小'
            
        # 格式化值
        if isinstance(value, str):
            formatted_value = f"'{value}'"
        else:
            formatted_value = str(value)
            
        lines.append(f"{' ' * indent}'{key}':{'' if len(key) >= 8 else ' ' * (8 - len(key))}{formatted_value},{comment}")
    lines.append(' ' * (indent - 4) + '}')
    return '\n'.join(lines)

# 更新trans_config.py文件
def update_trans_config(ms_configs, model_dir):
    config_file = "trans_config.py"
    
    if not os.path.exists(config_file):
        print(f"错误: 找不到 {config_file}")
        sys.exit(1)
    
    # 读取现有文件内容
    with open(config_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 获取当前时间
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # 移除现有的顶部注释（如果存在）
    if content.startswith('#'):
        # 找到第一个非注释行
        lines = content.split('\n')
        first_code_line = 0
        for i, line in enumerate(lines):
            if line.strip() and not line.strip().startswith('#'):
                first_code_line = i
                break
        content = '\n'.join(lines[first_code_line:])
    
    # 添加新的顶部注释
    header_comment = f"""# 配置文件自动生成
# 模型目录: {model_dir}
# 修改时间: {current_time}
"""
    updated_content = header_comment + content
    
    for config_name, config_dict in ms_configs.items():
        # 构建配置字符串
        config_str = dict_to_formatted_string(config_dict)
        
        # 替换文件中的配置块
        pattern = f"{config_name}=\{{[^{{}}]*\}}"
        replacement = f"{config_name}={config_str}"
        
        # 使用正则表达式替换
        updated_content = re.sub(pattern, replacement, updated_content, flags=re.DOTALL)
    
    # 写入更新后的内容
    with open(config_file, 'w', encoding='utf-8') as f:
        f.write(updated_content)
    
    print(f"成功更新 {config_file}")
    # 显示更新的配置
    for config_name, config_dict in ms_configs.items():
        print(f"\n{config_name}:")
        for key, value in config_dict.items():
            print(f"  {key}: {value}")

def main():
    print("=== 配置恢复工具 ===")
    # 获取模型目录名
    model_dir = get_model_dir_from_input()
    print(f"正在处理模型目录: {model_dir}")
    
    # 构建配置文件路径
    config_path = build_config_path(model_dir)
    print(f"配置文件路径: {config_path}")
    
    # 读取配置文件
    configs = read_training_configs(config_path)
    
    # 检查是否读取到配置
    if not configs:
        print("警告: 无法读取配置文件或配置文件为空")
        # 询问用户是否继续
        choice = input("配置文件读取失败，可能是旧版本模型目录。是否继续使用默认配置？(y/n): ")
        if choice.lower() != 'y':
            print("程序已退出")
            sys.exit(0)
        print("将使用默认配置继续")
    else:
        print(f"成功读取配置文件，包含 {len(configs)} 个配置项")
    
    # 提取MS配置
    ms_configs = extract_ms_configs(configs)
    
    # 检查是否找到MS配置
    if not ms_configs:
        print("警告: 未找到MS_CSV, MS_NML, MS_AVG, MS_END配置")
        # 尝试将整个配置作为一个阶段
        if configs:
            ms_configs = {'MS_CSV': configs}
        else:
            print("将使用默认配置")
            # 这里可以设置默认配置，或者让update_trans_config使用文件中已有的配置
    else:
        print(f"找到 {len(ms_configs)} 个MS配置")
    
    # 检查是否有on_train_epoch_start_source和init_params
    if 'on_train_epoch_start_source' in configs:
        print("找到on_train_epoch_start函数源码")
    else:
        print("注意: 未找到on_train_epoch_start函数源码(可能是旧版本模型)")
    
    if 'init_params' in configs:
        print("找到init_params配置")
    else:
        print("注意: 未找到init_params配置(可能是旧版本模型)")
    
    # 更新trans_config.py，传入model_dir参数用于生成注释
    update_trans_config(ms_configs, model_dir)
    
    print("\n配置恢复完成！")
    if not configs or not ms_configs:
        print("注意: 由于配置不完整，建议检查trans_config.py文件内容是否符合预期")

if __name__ == "__main__":
    main()